# -*- coding: utf-8 -*-

import pickle
from sklearn.naive_bayes import MultinomialNB
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from imblearn.under_sampling import OneSidedSelection
from imblearn.over_sampling import ADASYN 
from sklearn.metrics import f1_score
from sklearn.neighbors import KNeighborsClassifier

def solve_imb(X, y):
    # 过采样、欠采样
    oss = OneSidedSelection(random_state=0)
    X_res, y_res = oss.fit_resample(X, y)
    ada = ADASYN(random_state=42)
    X_res, y_res = ada.fit_resample(X_res, y_res)
    return X_res, y_res

def train_model(X,y, estimator, split=0.3):

    # 拆分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, 
                                                        y, 
                                                        test_size=split)
    # 训练模型
    model = estimator
    model.fit(X_train, y_train)
    
    # 评价模型
    y_train_pred = model.predict(X_train)
    print('测试集占比为', split)
    print('训练集中的：f1: ', f1_score(y_train, y_train_pred))

    y_test_pred = model.predict(X_test)
    print('测试集中的: f1: ', f1_score(y_test, y_test_pred))
    
    return model
    

if __name__ == '__main__':
    # 加载数据
    X = pickle.load(open(r'../附件/data_final.pkl', 'rb'))
    y = pickle.load(open(r'../附件/y.pkl', 'rb'))
    
    X_res, y_res = solve_imb(X, y)
    # 储存数据
    pickle.dump(X_res, open(r'../附件/X_res.pkl', 'wb'))
    pickle.dump(y_res, open(r'../附件/y_res.pkl', 'wb'))
    
    #X_res = pickle.load(open(r'../附件/X_res.pkl', 'rb'))
    #y_res = pickle.load(open(r'../附件/y_res.pkl', 'rb'))
    
    print('过采样、欠采样后数据： ', X_res.shape)
    # 定义模型
    model = LogisticRegression()
    # 训练模型，输出 F1：
    model = train_model(X_res, y_res, model, split=0.3)
    # 保存模型
    path_save = r'../附件/model.pkl'
    pickle.dump(model, open(path_save, 'wb'))
    
    
    
